{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X0EIJzvtglHw"
      },
      "source": [
        "# Nebius AI LLM\n",
        "This notebook shows how to use `Nebius AI` as an LLM. Nebius AI provides access to many state-of-the-art LLM models. Check out the full list of models [here](https://studio.nebius.ai/).\n",
        "\n",
        "Visit https://studio.nebius.ai/ and sign up to get an API key."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "an865BrVjQ7P"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h39hnQBdgVUt"
      },
      "outputs": [],
      "source": [
        "%pip install llama-index-llms-nebius"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0axqtiQDjS0O"
      },
      "source": [
        "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "TOYs4T2ngrxx"
      },
      "outputs": [],
      "source": [
        "!pip install llama-index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "3Y99wGYHg3j9"
      },
      "outputs": [],
      "source": [
        "from llama_index.llms.nebius import NebiusLLM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eNvx6QMlhG8Y"
      },
      "outputs": [],
      "source": [
        "# set api key in env or in llm\n",
        "# import os\n",
        "# os.environ[\"NEBIUS_API_KEY\"] = \"your api key\"\n",
        "\n",
        "llm = NebiusLLM(\n",
        "    model=\"mistralai/Mixtral-8x7B-Instruct-v0.1\", api_key=\"your_api_key\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "gFY9oqg-hTHR"
      },
      "outputs": [],
      "source": [
        "resp = llm.complete(\"Who is Paul Graham?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "V7tcT2PQhUYf",
        "outputId": "f9006a7a-b30e-40f9-9e36-8a1396832830"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " Paul Graham is a British computer scientist, venture capitalist, and essayist. He is best known as the co-founder of Y Combinator, a startup accelerator that has funded and helped grow many successful tech companies, including Dropbox, Airbnb, and Reddit.\n",
            "\n",
            "Before founding Y Combinator, Graham was a successful entrepreneur himself, having co-founded Viaweb, one of the first web-based application service providers, which was acquired by Yahoo in 1998.\n",
            "\n",
            "Graham is also known for his essays on startups, programming, and entrepreneurship, which have been widely read and influential in the tech community. He has written several books, including \"Hackers & Painters: Big Ideas from the Computer Age,\" \"On Lisp,\" and \"ANSI Common Lisp.\"\n",
            "\n",
            "In addition to his work in the tech industry, Graham has a strong interest in art and has exhibited his photography and painting in galleries. He has a degree in philosophy from Cornell University and a Ph.D. in computer science from Harvard University.\n"
          ]
        }
      ],
      "source": [
        "print(resp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kq_JzOlUjiqL"
      },
      "source": [
        "#### Call `chat` with a list of messages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "chEgCFCUjnkt"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.llms import ChatMessage\n",
        "\n",
        "messages = [\n",
        "    ChatMessage(\n",
        "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
        "    ),\n",
        "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
        "]\n",
        "resp = llm.chat(messages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kpgdwbz3jwbs",
        "outputId": "b020403c-5723-4fd5-f777-3e1895a24eca"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "assistant:  Arr matey! I'm Captain Redbeard, the fiercest pirate on the seven seas! My crew and I sail the waters in search of treasure and adventure. With me colorful personality and fearsome reputation, no one dares to cross us! Yarrr!\n"
          ]
        }
      ],
      "source": [
        "print(resp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mzg712SNj4rh"
      },
      "source": [
        "### Streaming"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NfAS7Owdj5cK"
      },
      "source": [
        "Using `stream_complete` endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "id": "ptqVIJF8j-QD"
      },
      "outputs": [],
      "source": [
        "response = llm.stream_complete(\"Who is Paul Graham?\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "mdFHUDmGkTlu",
        "outputId": "44450c63-b891-43c9-9006-c8db9d5093a8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " Paul Graham is a British computer scientist, venture capitalist, and essayist. He is best known as the co-founder of Y Combinator, a startup accelerator that has funded and helped grow many successful tech companies, including Dropbox, Airbnb, and Reddit.\n",
            "\n",
            "Before founding Y Combinator, Graham was a successful entrepreneur himself, having co-founded Viaweb, one of the first web-based application service providers, which was acquired by Yahoo in 1998.\n",
            "\n",
            "Graham is also known for his essays on startups, programming, and entrepreneurship, which have been widely read and influential in the tech community. He has written several books, including \"Hackers & Painters: Big Ideas from the Computer Age\" and \"The Hundred-Year Library\".\n",
            "\n",
            "In addition to his work in technology and venture capital, Graham is also a philanthropist and has made significant contributions to education and other causes."
          ]
        }
      ],
      "source": [
        "for r in response:\n",
        "    print(r.delta, end=\"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vkXzkg5Fkb8A"
      },
      "source": [
        "Using `stream_chat` endpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "iZIUR1gxkbHS"
      },
      "outputs": [],
      "source": [
        "from llama_index.core.llms import ChatMessage\n",
        "\n",
        "messages = [\n",
        "    ChatMessage(\n",
        "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
        "    ),\n",
        "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
        "]\n",
        "resp = llm.stream_chat(messages)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3-GuCPLrko-c",
        "outputId": "562412e6-98d4-482a-ee00-a17a643cdfd4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            " Arr matey! I'm Captain Redbeard, the fiercest pirate on the seven seas! My crew and I sail the waters in search of treasure and adventure. With me colorful personality and fearsome reputation, no one dares to cross us! Yarrr!"
          ]
        }
      ],
      "source": [
        "for r in resp:\n",
        "    print(r.delta, end=\"\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
